package lab.chapter;

import java.util.*;
import java.util.List;

/**
 * LL 算法实现的语法解析器
 * 包括：
 * 1. 计算 First 集合
 * 2. 计算 Follow 集合
 * 3. 通过递归下降算法，使用 First 和 Follow 集合实现带有预测能力的语法分析
 * 本例中，计算 First、Follow 集合，都是基于 GrammarNode 来做的。
 */
public class LLParser {

    public static void main(String[] args) {
        String script1 = "2+3*(4+5)";
        parse(script1, GrammarSample.expressionGrammar());

        String script2 = "int a = 0; a > 3;";
        parse(script2, GrammarSample.statementGrammar());

        String script3 = "int a = 0; if (a<3) a = 3;";
        parse(script3, GrammarSample.statementGrammar());
    }

    public static ASTNode parse(String script, GrammarNode grammarNode) {
        grammarNode.dump();

        // first 集合
        Map<GrammarNode, Set<String>> firstSets = calcFirstSets(grammarNode);
        System.out.println("\nFirst:");
        dumpFirstFollowSets(firstSets);
        // follow 集合
        Map<GrammarNode, Set<String>> followSets = calcFollowSets(grammarNode, firstSets);
        System.out.println("\nFollow:");
        dumpFirstFollowSets(followSets);

        // 词法分析
        List<Token> tokens = Lexer.tokenize(script);
        TokenReader tokenReader = new TokenReader(tokens);

        // 语法分析
        ASTNode rootNode = match(grammarNode, tokenReader, firstSets, followSets);
        if (rootNode != null && rootNode != ASTNode.EpsilonNode) {
            rootNode.dump();
            return rootNode;
        }
        return null;
    }

    /**
     * 计算first集合。采用不动点法
     * @param node 入口语法结点
     * @return map
     */
    public static Map<GrammarNode, Set<String>> calcFirstSets(GrammarNode node) {
        Map<GrammarNode, Set<String>> firstSets = new HashMap<>();
        boolean stable = calcFirstSets(node, firstSets, new HashSet<GrammarNode>());
        int i = 1;
        System.out.println("calculate first sets round: " + i++);
        while (!stable) {
            System.out.println("calculate first sets round: " + i++);
            stable = calcFirstSets(node, firstSets, new HashSet<GrammarNode>());
        }
        return firstSets;
    }

    /**
     * 对 first 集合做一次运算
     * @param node 入口节点
     * @param firstSets first 集合
     * @param calculated 已经计算的集合
     * @return 如果本次计算，first集合成员没有变动，则返回true。
     */
    private static boolean calcFirstSets(
            GrammarNode node, Map<GrammarNode,
            Set<String>> firstSets,
            Set<GrammarNode> calculated) {
        // 标记节点，避免重复调用
        calculated.add(node);
        Set<String> firstSet;

        // 获取、创建集合
        if(firstSets.containsKey(node)) {
            firstSet = firstSets.get(node);
        } else {
            firstSet = new HashSet<>();
            firstSets.put(node, firstSet);
        }

        if(node.isLeaf()) {
            return true;
        }
        boolean stable = true;
        // 把所有的子节点计算一遍
        for (GrammarNode child : node.children()) {
            if (!child.isLeaf() && !calculated.contains(child)) {
                // TODO 是否需要记录 Stable
                calcFirstSets(child, firstSets, calculated);
            }
        }

        List<GrammarNode> childToAdd = new LinkedList<>();
        if (node.getType() == GrammarNodeType.And) {
            childToAdd.add(node.getChild(0));

            // 一直找到一个不产生 epsilon 的符号
            for (GrammarNode child : node.children()) {
                childToAdd.add(child);
                if (!child.isNullable()) {
                    break;
                }
            }
        } else if (node.getType() == GrammarNodeType.Or) {
            // 加上所有的子节点
            childToAdd.addAll(node.children());
        }

        for (GrammarNode child : childToAdd) {
            if (!child.isLeaf()) {
                Set<String> childSet = firstSets.get(child);
                if (!firstSet.containsAll(childSet)) {
                    firstSet.addAll(childSet);
                    stable = false;
                }
            } else if (child.isToken()) {
                if (!firstSet.contains(child.getToken().getType())) {
                    stable = false;
                    firstSet.add(child.getToken().getType());
                }
            }

            if (child.isNullable()) {
                firstSet.add("EPSILON");
            }
        }
        return stable;
    }

    /**
     * 对所有节点计算 follow 集合
     * @param node 入口节点
     * @param firstSets first 集合
     * @return follow 集合
     */
    public static Map<GrammarNode, Set<String>> calcFollowSets(GrammarNode node,
                                                               Map<GrammarNode, Set<String>> firstSets) {
        Map<GrammarNode, Set<String>> followSets = new HashMap<>();
        Map<GrammarNode, Set<GrammarNode>> rightChildrenSets = new HashMap<>();
        // 不动点计算 follow 集合
        int i = 1;
        System.out.println("follow set round: " + i++);
        boolean stable = calcFollowSets(
                node, followSets, rightChildrenSets, firstSets, new HashSet<GrammarNode>());

        while (!stable) {
            System.out.println("follow set round: " + i++);
            stable = calcFollowSets(node, followSets, rightChildrenSets, firstSets, new HashSet<GrammarNode>());
        }

        return followSets;
    }

    /**
     * 计算一遍 follow 节点
     * @param node 入口节点
     * @param followSets follow sets
     * @param rightChildrenSets 右节点集合
     * @param firstSets first 集合
     * @param calculated 已经计算的节点
     * @return bool
     */
    private static boolean calcFollowSets(GrammarNode node,
                                          Map<GrammarNode, Set<String>> followSets,
                                          Map<GrammarNode, Set<GrammarNode>> rightChildrenSets,
                                          Map<GrammarNode, Set<String>> firstSets,
                                          Set<GrammarNode> calculated) {
        calculated.add(node);
        if (node.isLeaf()) {
            return true;
        }
        boolean stable = true;
        // 等价于：if (rightChildren == null) {rightChildren = new HashSet<GrammarNode>(); rightChildrenSets.put(node, rightChildren);}
        Set<GrammarNode> rightChildren = rightChildrenSets.computeIfAbsent(node, k -> new HashSet<>());
        if (node.getType() == GrammarNodeType.And) {
            for (int i = 0; i < node.getChildCount(); i++) {
                GrammarNode left = node.getChild(i);
                if (!left.isLeaf()) {
                    if (!calculated.contains(left)) {
                        if (!calcFollowSets(left, followSets, rightChildrenSets, firstSets, calculated)) {
                            stable = false;
                        }
                    }

                    Set<String> followSet = followSets.computeIfAbsent(left,  k -> new HashSet<>());
                    if (i == node.getChildCount()) {
                        rightChildren.add(left);
                    } else {
                        boolean foundNotNull = false;
                        for (int j = i + 1; j < node.getChildCount(); j++){
                            GrammarNode right = node.getChild(j);
                            Set<String> tempFollowSet = new HashSet<>();
                            if (!right.isLeaf()) {
                                if (firstSets.get(right) == null) {
                                    System.out.println("");
                                }
                                tempFollowSet.addAll(firstSets.get(right));
                            } else if (right.isToken()) {
                                tempFollowSet.add(right.getToken().getType());
                            }
                            assert (tempFollowSet.size() > 0);
                            if (!followSet.containsAll(tempFollowSet)) {
                                followSet.addAll(tempFollowSet);
                                stable = false;
                            }
                            if (!addToRightChild(left, tempFollowSet, followSets, rightChildrenSets, new HashSet<GrammarNode>())){
                                stable = false;
                            }

                            if (!right.isNullable()) {
                                // 找到一个非空的
                                foundNotNull = true;
                                break;
                            }
                        }
                        // 本节点也是最右节点
                        if (!foundNotNull) {
                            rightChildren.add(left);
                        }
                    }
                }
            }
        } else if (node.getType() == GrammarNodeType.Or) {
            for (GrammarNode child : node.children()) {
                if (!child.isLeaf()) {
                    rightChildren.add(child);
                    if (!calculated.contains(child)) {
                        // 再计算一次节点
                        if (!calcFollowSets(child, followSets, rightChildrenSets, firstSets, calculated)) {
                            stable = false;
                        }
                    }
                }
            }
        }

        return stable;
    }

    /**
     * 把某个节点的 Follow 集合，也给它所有右节点分支后代节点
     * @param node
     * @param followSet
     * @param followSets
     * @param rightChildrenSets
     * @param added
     * @return
     */
    private static boolean addToRightChild(GrammarNode node,
                                           Set<String> followSet,
                                           Map<GrammarNode, Set<String>> followSets,
                                           Map<GrammarNode, Set<GrammarNode>> rightChildrenSets,
                                           Set<GrammarNode> added) {
        added.add(node);
        boolean stable = true;
        Set<GrammarNode> rightChildren = rightChildrenSets.get(node);
        for (GrammarNode rightChild : rightChildren) {
            if (!rightChild.isLeaf() && !added.contains(rightChild)) {
                // 入演算
                Set<String> childFollowSet = followSets.computeIfAbsent(rightChild, k -> new HashSet<>());
                if (!childFollowSet.containsAll(followSet)) {
                    childFollowSet.addAll(followSet);
                    stable = false;
                }
                if (!addToRightChild(rightChild, followSet, followSets, rightChildrenSets, added)) {
                    stable = false;
                }
            }
        }
        return stable;
    }

    private static ASTNode match(GrammarNode grammarNode,
                                 TokenReader tokenReader,
                                 Map<GrammarNode, Set<String>> firstSets,
                                 Map<GrammarNode, Set<String>> followSets) {
        ASTNode node = new ASTNode(grammarNode.getName());
        if (grammarNode.getChildCount() > 0) {
            assert (grammarNode.getType() == GrammarNodeType.And || grammarNode.getType() == GrammarNodeType.Or);
            // 对于and类型，每个子节点要依次全部匹配
            if (grammarNode.getType() == GrammarNodeType.And) {
                for (GrammarNode child : grammarNode.children()) {
                    ASTNode childNode = match(child, tokenReader, firstSets, followSets);
                    if (childNode != null) {
                        // 子节点匹配成功
                        if (childNode != ASTNode.EpsilonNode) {
                            node.addChild(childNode);
                        }
                    } else {
                        // 子节点匹配失败，回溯，报错
                        System.out.println("failed to match: " + child);
                        break;
                    }
                }
            } else {
                // 对于 or 类型，通过预测，知道采用哪个产生式
                boolean matched = false;
                Token token = tokenReader.peek();
                for (GrammarNode child : grammarNode.children()) {
                    if (child.isToken()) {
                        if (token.getType().equals(child.getToken().getType())) {
                            matched = true;
                        }
                    } else if (child.getChildCount() > 0) {
                        Set<String> firstSet = firstSets.get(child);
                        assert firstSet != null;
                        if(firstSet.contains(token.getType())) {
                            matched = true;
                        }
                    }
                    if (matched) {
                        ASTNode childNode = match(child, tokenReader, firstSets, followSets);
                        if(childNode != null) {
                            if (childNode != ASTNode.EpsilonNode) {
                                node.addChild(childNode);
                            }
                        } else {
                            System.out.println("failed to match: " + child);
                        }
                        break;
                    }
                }
                if (!matched) {
                    // 看看是否要产生 Epsilon
                    boolean epsilon = false;
                    if (grammarNode.isNullable()) {
                        Set<String> followSet = followSets.get(grammarNode);
                        assert followSet != null;
                        if (followSet.contains(token.getType())) {
                            node = ASTNode.EpsilonNode;
                            epsilon = true;
                        }
                    }
                    if (!epsilon) {
                        System.out.println("unable to find a selection for: " + grammarNode);
                    }
                }
            }
        } else if (grammarNode.getType() == GrammarNodeType.Epsilon) {
            node = ASTNode.EpsilonNode;
        } else if (grammarNode.isToken()) {
            // 叶子节点，看看自身的类型跟语法要求的是否一致
            Token token = tokenReader.peek();
            if (token != null) {
                assert grammarNode.getToken() != null;
                if (token.getType().equals(grammarNode.getToken().getType())) {
                    // 读取下一个
                    tokenReader.read();
                    node = new ASTNode((token.getType()));
                    node.setText(token.getText());
                }
                // 匹配失败
            }
        } else {
            System.out.println("unsupported grammarNode: " + grammarNode);
        }

        // 如果全部子节点返回的都是 Epsilon，自身也置为Epsilon
        if (grammarNode.getChildCount() > 0 && node.getChildCount() == 0) {
            node = ASTNode.EpsilonNode;
        }
        return node;
    }

    /**
     * 打印 First or Follow 集合
     * @param sets 待打印的集合
     */
    private static void dumpFirstFollowSets(Map<GrammarNode, Set<String>> sets) {
        for (GrammarNode node : sets.keySet()) {
            Set<String> set = sets.get(node);
            StringBuilder s = new StringBuilder(node.toString() + " : ");
            for (String tokenName : set) {
                s.append(" ").append(tokenName);
            }
            System.out.println(s.toString());
        }
    }
}
